# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

from tornado import gen
from coin.exchange.okex_v3.ws.private_subscriber import (OkexWsPrivateClientBase, OperationType)


class OkexFuturesWsPrivateClient(OkexWsPrivateClientBase):
  def subscribe_futures_balance(self, currency):
    channel = 'futures/account:' + currency
    return self.send_request(OperationType.SUBSCRIBE, [channel])

  def subscribe_futures_positions(self, product):
    channel = 'futures/position:' + product
    return self.send_request(OperationType.SUBSCRIBE, [channel])

  def subscribe_futures_order(self, product):
    channel = 'futures/order:' + product
    return self.send_request(OperationType.SUBSCRIBE, [channel])

  @gen.coroutine
  def prepare_futures_subscription(self, products):
    for product in products:
      response = yield self.subscribe_futures_positions(product)
      event = response['event']
      if event != 'error':
        self._logger.info('Subscribe order response: %s\n', response)
        channel = response.get('channel')
        self._subscribed_channels.add(channel)
      else:
        self._logger.info('Position subscription failed: %s', response)
        self.stop()
        return

      response = yield self.subscribe_futures_order(product)
      if event != 'error':
        self._logger.info('Subscribe order response: %s\n', response)
        channel = response.get('channel')
        self._subscribed_channels.add(channel)
      else:
        self._logger.info('Trade subscription failed: %s', response)
        self.stop()
        return
